from template import *
from collect import collect

import sys

itime = 0
if len(sys.argv) == 2:
	t = int(sys.argv[1])
elif len(sys.argv) > 2:
	print("error: should have one argument, as time step")
else:
	t = itime


##inite the fig of matplotlib
fig = plt.figure(figsize=(8,8))
fig.subplots_adjust(top=0.9,bottom=0.1,wspace=0.5,hspace=0.55)

label = ['1', '2', '3', '4', '5', '6', '7']
#============ vdf parallel ======================================================
ax = fig.add_subplot(3,1,1)

val = collect("/Diagnostic/", "vdf_velocity_array")
x = val[t, 0, 1, :]

val = collect("/Diagnostic/", "vdf_parallel_to_B")
y = val[t, :, 1, :]

for i in range(0, 7):
    line = ax.plot(x, y[i, :], label = label[i])


ax.legend(loc = 1, framealpha=1)

#============ vdf perpendicular ======================================================
ax = fig.add_subplot(3,1,2)

val = collect("/Diagnostic/", "vdf_velocity_array")
x = val[t, 0, 1, :]

val = collect("/Diagnostic/", "vdf_perpendicular_to_B")
y = val[t, :, 1, :]

for i in range(0, 7):
    line = ax.plot(x, y[i, :], label = label[i])

ax.legend(loc = 1, framealpha=1)

#============ edf ======================================================
ax = fig.add_subplot(3,1,3)

val = collect("/Diagnostic/", "edf_energy_array")
x = val[t, 0, 1, :]

val = collect("/Diagnostic/", "edf")
y = val[t, :, 1, :]

for i in range(0, 7):
    line = ax.plot(x, y[i, :], label = label[i])

ax.legend(loc = 1, framealpha=1)

fig_file_name = "ion_vdf" + str(t) + ".png"
fig.savefig(fig_file_name, dpi = 300)

